import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
import torchvision
from torchvision import transforms
from torch.utils.data import DataLoader

# print(torch.__version__)


# 对数据做归一化 (-1,1)
transform = transforms.Compose([
    transforms.ToTensor(),  # (0,1)
    transforms.Normalize(0.5, 0.5),  # (-1,1)
])

train_ds = torchvision.datasets.MNIST(root='./dataset', train=True, transform=transform, download=True)

dataloader = DataLoader(train_ds, batch_size=64, shuffle=True)

imgs, _ = next(iter(dataloader))
print(imgs.shape)


class Generate(nn.Module):
    def __init__(self):
        super(Generate, self).__init__()
        self.main = nn.Sequential(
            nn.Linear(100, 256),
            nn.ReLU(),
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512, 28 * 28),
            nn.Tanh(),
        )

    def forward(self, x):
        img = self.main(x)
        img = img.view(-1, 28, 28, 1)
        return img


class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.main = nn.Sequential(
            nn.Linear(28 * 28, 512),
            nn.LeakyReLU(),
            nn.Linear(512, 256),
            nn.LeakyReLU(),
            nn.Linear(256, 1),
            nn.Sigmoid(),
        )

    def forward(self, x):
        x = x.view(-1, 28 * 28)
        x = self.main(x)
        return x


device = 'cuda' if torch.cuda.is_available() else 'cpu'

G = Generate().to(device)
D = Discriminator().to(device)

D_optim = torch.optim.Adam(D.parameters(), lr=0.0001)
G_optim = torch.optim.Adam(G.parameters(), lr=0.0001)

loss_fn = torch.nn.BCELoss()


def gen_img_plot(model, test_input):
    prediction = np.squeeze(model(test_input).detach().cpu().numpy())
    fig = plt.figure(figsize=(4, 4))
    for i in range(16):
        plt.subplot(4, 4, i + 1)
        plt.imshow((prediction[i] + 1) / 2)
        plt.axis('off')
    plt.show()


test_input = torch.randn(16, 100, device=device)

D_loss = []
G_loss = []

for epoch in range(20):
    d_epoch_loss = 0
    g_epoch_loss = 0
    count = len(dataloader)
    for step, (img, _) in enumerate(dataloader):
        img = img.to(device)
        size = img.size(0)
        random_noise = torch.randn(size, 100, device=device)

        D_optim.zero_grad()
        real_output = D(img)
        d_real_loss = loss_fn(real_output, torch.ones_like(real_output))
        d_real_loss.backward()

        gen_img = G(random_noise)
        fake_output = D(gen_img.detach())
        d_fake_loss = loss_fn(fake_output, torch.zeros_like(fake_output))
        d_fake_loss.backward()

        d_loss = d_real_loss + d_fake_loss
        D_optim.step()

        G_optim.zero_grad()
        fake_output = D(gen_img)
        g_loss = loss_fn(fake_output, torch.ones_like(fake_output))
        g_loss.backward()
        G_optim.step()

        with torch.no_grad():
            d_epoch_loss += d_loss
            g_epoch_loss += g_loss

    with torch.no_grad():
        d_epoch_loss /= count
        g_epoch_loss /= count
        D_loss.append(d_epoch_loss)
        G_loss.append(g_epoch_loss)
        print('Epoch:', epoch)
        gen_img_plot(G, test_input)